package com.yujiahao;

import org.apache.flink.kinesis.shaded.com.amazonaws.ClientConfiguration;
import org.apache.flink.kinesis.shaded.com.amazonaws.auth.DefaultAWSCredentialsProviderChain;
import org.apache.flink.kinesis.shaded.com.amazonaws.retry.PredefinedRetryPolicies;
import org.apache.flink.kinesis.shaded.com.amazonaws.services.kinesis.AmazonKinesis;
import org.apache.flink.kinesis.shaded.com.amazonaws.services.kinesis.AmazonKinesisClientBuilder;
import org.apache.flink.kinesis.shaded.com.amazonaws.services.kinesis.clientlibrary.interfaces.v2.IRecordProcessor;
import org.apache.flink.kinesis.shaded.com.amazonaws.services.kinesis.clientlibrary.types.InitializationInput;
import org.apache.flink.kinesis.shaded.com.amazonaws.services.kinesis.clientlibrary.types.ProcessRecordsInput;
import org.apache.flink.kinesis.shaded.com.amazonaws.services.kinesis.clientlibrary.types.ShutdownInput;
import org.apache.flink.kinesis.shaded.com.amazonaws.services.kinesis.model.DescribeStreamRequest;
import org.apache.flink.kinesis.shaded.com.amazonaws.services.kinesis.model.Record;
import software.amazon.awssdk.core.internal.retry.RetryHandler;

import java.nio.charset.StandardCharsets;
import java.util.List;

/**
 * @Description: 使用IRecordProcessor接口 ( package com.amazonaws.services.kinesis.clientlibrary.interfaces.v2)
 * @Author: yujiahao
 * @Date: 2023/8/2 17:33
 * @Version: 1.0
 */

public class KinesisConsumeIRecordProcessorV2 implements IRecordProcessor {

    private static AmazonKinesis client;
    private String streamName = "test-kds-api";
    private String sequenceNumberOfPreviousRecord;
    DescribeStreamRequest describeStreamRequest;

    @Override
    public void initialize(InitializationInput initializationInput) {
        // 初始化 AmazonKinesis 客户端
        ClientConfiguration clientConfiguration = new ClientConfiguration();
        clientConfiguration.setMaxConnections(50); // 设置连接池最大连接数
        clientConfiguration.setConnectionTimeout(5000); // 设置连接超时时间，单位：毫秒

        AmazonKinesisClientBuilder clientBuilder = AmazonKinesisClientBuilder.standard()
                .withCredentials(new DefaultAWSCredentialsProviderChain())
                .withClientConfiguration(clientConfiguration);
        client = clientBuilder.build();

    }

    @Override
    public void processRecords(ProcessRecordsInput processRecordsInput) {
        List<Record> records = processRecordsInput.getRecords();

        // 处理从 Kinesis 数据流中获取的记录
        for (Record record : records) {
            String data = new String(record.getData().array(), StandardCharsets.UTF_8);
            String partitionKey = record.getPartitionKey();
            String sequenceNumber = record.getSequenceNumber();

            // 在这里可以对数据进行进一步处理或保存到数据库等操作
            System.out.println("Received data: " + data);
            System.out.println("Partition key: " + partitionKey);
            System.out.println("Sequence number: " + sequenceNumber);

            // 在此处可以将当前记录的序列号设置为下一次处理的上一条记录的序列号
            sequenceNumberOfPreviousRecord = sequenceNumber;
        }
    }

    @Override
    public void shutdown(ShutdownInput shutdownInput) {
        // 关闭 AmazonKinesis 客户端
        client.shutdown();
    }
}
